'''自制的aria2下载服务器管理,其实也只有检查连接状态与添加下载表两个功能.
  使用时就只用addUri方法就行了,简单.
  
  提交批量下载还要批量重命名,貌似暂时找不到办法.
  
  By setycyas @2019/9/7
'''

import time
import ssl
import json
import sys

import requests

ssl._create_default_https_context = ssl._create_unverified_context

## 用端口与记录文件进行初始化,记录文件为空则不进行记录
class Aria2DownloadServer:

  def __init__(self, port = '6800'):
    self._rpc = r'http://localhost:{0}/jsonrpc'.format(port) # 访问aria2的接口
    self._isLinked = self.checkLink() # 连接是否成功
    return
    
  ## 测试连接是否成功
  def checkLink(self):
    timeStamp = int(time.time()*1000)
    url = self._rpc+'?tm=1567832001167'
    ## 制作表单数据
    data = {
      'jsonrpc':'2.0',
      'method':'aria2.tellActive',
      'id':1
    }
    ## 开始查询,成功连接服务器返回True,否则返回False
    try:
      resp = requests.post(url, json = data)
      id = json.loads(resp.text)['id']
      if id:
        print('连接Aria2服务器成功!')
        return True
      else:
        print('Aria2服务器无法识别请求!')
        return False
    except:
      print('无法连接Aria2服务器!')
      return False
      
  ## 添加下载文件,添加成功返回True.uri是下载地址
  ## out:下载的文件名,可以留空
  ## dir:下载目录
  ## headers:header字符串表,如['referer:www.a2b.com','ua:firefox']
  ## 后面2个参数是分块数,单一服务器最大连接数
  def addUri(self, uri, out = None, dir = sys.path[0], headers = [],
    split = 1, max_connection_per_server = 16):
    ## 如果连接失败,终止
    if not self._isLinked:
      print('Aria2服务器未准备就绪,不能添加下载!')
      return False
    ## 开始添加下载
    timeStamp = int(time.time()*1000)
    url = self._rpc+'?tm='+str(timeStamp)
    ## 制作表单数据
    data = [{
      'jsonrpc':'2.0',
      'method':'aria2.addUri',
      'id':1,
      'params':[[uri], {
        'dir': dir, # 下载目录
        'split': str(split), # 分块数
        'max-connection-per-server': str(max_connection_per_server)
      }]
    }]
    if headers:
      data[0]['params'][1]['header'] = ('\n').join(headers)
    if out:
      data[0]['params'][1]['out'] = out
    ## 提交请求
    try:
      resp = requests.post(url, json = data)
      id = json.loads(resp.text)[0]['id']
      if id:
        return True
      else:
        return False
    except:
      print('无法连接Aria2服务器!')
      return False
      
    
if __name__ == '__main__':
  aria2 = Aria2DownloadServer()
  uri = r'https://img.ku137.net/piccc/2019/img/190516/16031416-02261.jpg'
  headers = [r'Referer:https://img.ku137.net/']
  aria2.addUri(uri, out = '1.jpg', 
    dir = r'F:\Download\PyDownloadTest', 
    headers = headers
  )